%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Program: individual_cevs.m
% By: Stephie Fried, Kevin Novan, and William Peterman
% Date: Spring 2024
% Purpose: Calculates the ex-post CEVs
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
clc;
close all;
clear all;

bpath = 'C:\Users\l1sxf17\Dropbox (FRB SF)\Research\Will\equity_efficiency_project\Replication';
cd(bpath);

options = optimset('Algorithm','Levenberg-Marquardt','TolFun',10e-15,'TolX',10e-15, ...
    'MaxFunEvals', 10000, 'MaxIter', 10000, 'Display','none');

%% load data
load('Programs/Matlab/aggregates')

labor_inc_raw = importdata('Fortran_output/baseline/sim_labor_earnings_output.dat');
capital_inc_raw = importdata('Fortran_output/baseline/sim_capital_earnings_output.dat');
capital_tax_raw =importdata('Fortran_output/baseline/sim_capital_tax_output.dat');
labor_tax_raw =importdata('Fortran_output/baseline/sim_labor_tax_output.dat');
c_raw =importdata('Fortran_output/baseline/sim_consumption_output.dat');
ec_raw =importdata('Fortran_output/baseline/sim_energy_output.dat');
h_raw =importdata('Fortran_output/baseline/sim_labor_output.dat');

%% Find lifetime income  and welfare for each individual in the baseline
ec = zeros(81, 10000); %81 ages, 10000 people
c = zeros(81, 10000); %81 ages, 10000 people
labor_inc = zeros(81, 10000); 
labor_tax = zeros(81, 10000);
capital_inc = zeros(81, 10000); 
capital_tax =  zeros(81, 10000);
h =  zeros(81, 10000);

counter =0;
for i = 1:10000
    for j = 1:J %age j, individual i
        counter = counter+1;
        c(j, i) = c_raw(counter);
        ec(j, i) = ec_raw(counter);
        h(j, i) = h_raw(counter);
        labor_inc(j, i) =labor_inc_raw(counter);
        labor_tax(j, i) =labor_tax_raw(counter);
        capital_inc(j, i) =capital_inc_raw(counter);
        capital_tax(j, i) =capital_tax_raw(counter);
    end
end

life_inc = zeros(10000,1); 
for i = 1:10000
    life_inc(i) =labor_inc(1,i) + capital_inc(1,i) - labor_tax(1,i) - capital_tax(1,i); 
    for j = 2:1:J
        life_inc(i) = life_inc(i) + (1/(1+r*(1-lambdak)))^(j-1)*...
            (labor_inc(j,i) + capital_inc(j,i) - labor_tax(j,i) - capital_tax(j,i));
    end
end

cTilde = c.^gamma.*(ec - ebar).^(1-gamma); %ctilde
utility_base = cTilde.^(1-theta1)/(1-theta1) - chi*h.^(1+1/theta2)/(1+1/theta2); %U(ctilde)
mu = cTilde.^(-theta1);

welfare = zeros(10000,1);

for i = 1:10000
    for j =1:J
        if j == 1 
            welfare(i) = utility_base(j, i);
        else
            welfare(i) = beta^(j-1)*prod(survive(1:j-1))*utility_base(j,i) + welfare(i);
        end
    end
end

c_base = c; 
h_base = h; 
ec_base = ec; 
welfare_base = welfare;
labor_tax_base =labor_tax; 

%% Calculate "CEV" for each individual and for each policy. 
%optimal_porg is like "simple_prog", rebate only by increasing the
%progressivity of the labor tax
%Kevin experiment:  use the the income dependent rebate to achieve the
%samme redistribution as under the optimal policy
policy_vec = {'optimal', 'simple_capital', 'simple_labor', 'simple_lump', 'optimal_prog', 'MEID'};
folder_vec ={'optimal', 'simple_capital', 'simple_labor', 'simple_lump', 'prog', 'income_dependent'};
cevimat = zeros(10000, length(policy_vec));
weldiff=cevimat; 
fval = zeros(10000, length(policy_vec)); 
exit = zeros(10000, length(policy_vec)); 

for pi = 1:length(policy_vec)
    policy = policy_vec{pi}
    folder = folder_vec{pi};
   
    if (pi >1 & pi <5)
        c_raw =importdata(strcat('Fortran_output/', folder, '/sim_consumption_output.dat'));
        ec_raw =importdata(strcat('Fortran_output/', folder, '/sim_energy_output.dat'));
        h_raw =importdata(strcat('Fortran_output/', folder, '/sim_labor_output.dat'));
        k_raw = importdata(strcat('Fortran_output/', folder, '/sim_capital_output.dat'));
    elseif (pi ==5 || pi ==1)
        c_raw =importdata(strcat('Fortran_output/', folder, '/sim_consumption_output_optimal.dat'));
        ec_raw =importdata(strcat('Fortran_output/', folder, '/sim_energy_output_optimal.dat'));
        h_raw =importdata(strcat('Fortran_output/', folder, '/sim_labor_output_optimal.dat'));
        k_raw = importdata(strcat('Fortran_output/', folder, '/sim_capital_output_optimal.dat'));
    else
        c_raw =importdata(strcat('Fortran_output/', folder, '/sim_consumption_output_kevin_max.dat'));
        ec_raw =importdata(strcat('Fortran_output/', folder, '/sim_energy_output_kevin_max.dat'));
        h_raw =importdata(strcat('Fortran_output/', folder, '/sim_labor_output_kevin_max.dat'));
        k_raw = importdata(strcat('Fortran_output/', folder, '/sim_capital_output_kevin_max.dat'));
    end
  
    ec = zeros(81, 10000); %81 ages, 10000 people
    c = zeros(81, 10000); %81 ages, 10000 people
    h =  zeros(81, 10000);
    k = zeros(82, 10000);

    counter =0;
    for i = 1:10000
        for j = 1:J %age j, individual i
            counter = counter+1;
            c(j, i) = c_raw(counter);
            ec(j, i) = ec_raw(counter);
            h(j, i) = h_raw(counter);
            if (pi ==1)
                labor_tax_opt(j,i) = labor_tax_raw(counter); 
            end
        end
    end
    counter =0; 
    for i = 1:10000
        for j = 1:J+1 
            counter = counter +1;
            k(j, i) = k_raw(counter);
        end
    end

    cTilde = c.^gamma.*(ec - ebar).^(1-gamma); %ctilde
    utility = cTilde.^(1-theta1)/(1-theta1) - chi*h.^(1+1/theta2)/(1+1/theta2); %U(ctilde)
    
    welfare = zeros(10000,1);
    
    for i = 1:10000
        for j =1:J
            if j == 1
                welfare(i) = utility(j, i);
            else
                welfare(i) = beta^(j-1)*prod(survive(1:j-1))*utility(j,i) + welfare(i);
            end
        end
    end
    
      if (pi ==1)
        c_opt = c;
        ec_opt = ec;
        h_opt = h;
        k_opt =k;
        welfare_opt = welfare;
    elseif (pi ==2)
        c_sk = c;
        ec_sk = ec;
        h_sk = h;
        k_sk =k;
    elseif (pi ==3)
        c_sl = c;
        ec_sl = ec;
        h_sl = h;
        k_sl =k;
    elseif (pi ==4)
        c_lump = c;
        ec_lump = ec;
        h_lump = h;
        k_lump =k;
    elseif (pi ==5)
        c_prog = c;
        ec_prog = ec;
        h_prog = h;
        k_prog =k;
    else
        c_id = c;
        ec_id = ec;
        h_id = h;
        k_id =k;
        
      end


    cd('Programs/Matlab'); 
    for i = 1:10000
        target = welfare(i);
        guess = 0.3;
        [cevimat(i, pi),fval(i, pi),exit(i, pi)] = fsolve(@cevFunc, guess, options, c_base(:, i), ec_base(:, i), h_base(:, i), ...
            target, beta, chi, gamma, theta1, theta2, ebar, J, survive);
        weldiff(i, pi) = welfare(i) - welfare_base(i); 
    end
    cd(bpath); 


end
save ('Programs/Matlab/results_cevi')
%% Fractions better off and fractions that prefer the optimal

for  i = 1:6
    temp= length(find(cevimat(:,i)>0));
    frac_better_off(i) = temp/10000;
end

%policy_vec = {'optimal', 'simple_capital', 'simple_labor', 'simple_lump', 'optimal_prog', 'MEID'};
for i = 2:1:6
    temp = length(find(cevimat(:,1)> cevimat(:, i))); 
    frac_prefer(i) = temp/10000;
end

save('Programs/Matlab/frac_better_off', 'frac_better_off', 'frac_prefer');

